import socket
import time
import struct
from PIL import Image
from io import BytesIO

NIL = 0
ACK = 1
NAK = 2

VER = '2.0'

class RDT20Server:
    def __init__(self, addr : tuple[str, int]) -> None:
        self.serverSocket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self.serverSocket.bind(addr)
        self.data = bytearray()
        self.tmpdata = None
        self.target = None
        pass


    def recvpkt(self) -> bytes:
        data, target = self.serverSocket.recvfrom(1024)
        self.target = target
        return data
    
    def makepkt(self, rawdata: bytes, confirmation: int):
        # rdt 2.0 报文格式:
        # ack/nak 1byte
        # length 2byte
        # checksum 2byte
        pkt = struct.pack(
            "!BHH%ds" % len(rawdata), confirmation, len(rawdata), 0, rawdata
        )
        csum = self.checksum(pkt)
        return struct.pack(
            "!BHH%ds" % len(rawdata), confirmation, len(rawdata), csum, rawdata
        )

    def sendpkt(self, pkt: bytes):
        self.serverSocket.sendto(pkt, self.target)

    def checksum(self, data: bytes) -> int:
        csum = 0
        for i in range(1, len(data), 2):
            csum += int.from_bytes(data[i - 1 : i + 1], "big")
        if len(data) % 2 == 1:
            csum += data[-1] << 8 & 0xFF00
        csum = (csum >> 16) + (csum & 0xFFFF)
        csum = (csum >> 16) + (csum & 0xFFFF)
        csum = ~csum & 0xFFFF
        return csum >> 8 | ((csum << 8) & 0xFF00)
    
    def valid(self, pkt: bytes) -> bool:
        return self.checksum(pkt) == 0

    def extract(self, pkt: bytes) -> tuple[int, int, int, bytes]:
        confirmation, length, csum = struct.unpack("!BHH", pkt[:5])
        data = struct.unpack_from("!%ds" % length, pkt, 5)[0]
        return confirmation, length, csum, data

    def run(self):
        while True:
            pktrcvd = self.recvpkt()
            if self.valid(pktrcvd):
                confirmation, length, csum, self.tmpdata = self.extract(pktrcvd)
                if self.tmpdata == b'':
                    print('data fully received.')
                    self.sendpkt(self.makepkt(b'', ACK))
                    break
                print('data received at ', len(self.data))
                self.data += self.tmpdata
                self.sendpkt(self.makepkt(b'', ACK))
            else:
                print('data corrupted at ', len(self.data))
                self.sendpkt(self.makepkt(b'', NAK))
        pass

    pass

s = RDT20Server(('127.0.0.1', 65432))
s.run()

Image.open(BytesIO(s.data)).show()

open(VER + '-server.png', 'wb').write(s.data)
input("Press any key to exit...")
